#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from nav_msgs.msg import OccupancyGrid, Path, Odometry
from sensor_msgs.msg import LaserScan
from geometry_msgs.msg import Twist
import numpy as np
import cv2
import json

"""管理ROS通信"""
class ROSBridge(Node):
    def __init__(self, socketio):
        super().__init__('ros_web_bridge')
        self.socketio = socketio
        
        # 订阅器
        self.create_subscription(OccupancyGrid, '/map', self.map_callback, 10)
        self.create_subscription(Path, '/plan', self.path_callback, 10)
        self.create_subscription(Odometry, '/odom', self.odom_callback, 10)
        self.create_subscription(LaserScan, '/scan', self.scan_callback, 10)
        
        # 发布器
        self.cmd_vel_pub = self.create_publisher(Twist, '/cmd_vel', 10)
        
        # 数据存储
        self.map_data = None
        self.map_info = None
        self.robot_pose = None
        self.laser_data = None
        self.current_path = []
        
        # 服务客户端
        self.setup_services()

    def setup_services(self):
        from std_srvs.srv import Empty
        from nav2_msgs.srv import LoadMap
        
        # 地图服务
        self.load_map_client = self.create_client(
            LoadMap, '/nav2_navigation/load_map')
        
        # 生命周期服务
        self.startup_nav2 = self.create_client(
            Empty, '/nav2_navigation/startup')
        
        # 其他服务可根据需要添加

    def map_callback(self, msg):
        """处理地图数据"""
        try:
            # 转换地图数据
            np_map = np.array(msg.data, dtype=np.int8).reshape(
                (msg.info.height, msg.info.width))
            np_map = np.where(np_map == -1, 127, np_map * 255 // 100)
            
            # 编码为PNG
            _, buffer = cv2.imencode('.png', np_map)
            self.map_data = buffer.tobytes()
            
            # 存储地图元数据
            self.map_info = {
                'resolution': msg.info.resolution,
                'width': msg.info.width,
                'height': msg.info.height,
                'origin': {
                    'x': msg.info.origin.position.x,
                    'y': msg.info.origin.position.y
                }
            }
            
            # 通过SocketIO发送更新
            self.socketio.emit('map_update', {
                'info': self.map_info,
                'timestamp': self.get_clock().now().nanoseconds
            })
            
        except Exception as e:
            self.get_logger().error(f'Map processing error: {str(e)}')

    def path_callback(self, msg):
        """处理路径数据"""
        self.current_path = [
            {
                'x': pose.pose.position.x,
                'y': pose.pose.position.y,
                'theta': 2 * np.arctan2(
                    pose.pose.orientation.z,
                    pose.pose.orientation.w)
            } for pose in msg.poses
        ]
        self.socketio.emit('path_update', {
            'path': self.current_path,
            'timestamp': self.get_clock().now().nanoseconds
        })

    def odom_callback(self, msg):
        """处理里程计数据"""
        self.robot_pose = {
            'x': msg.pose.pose.position.x,
            'y': msg.pose.pose.position.y,
            'theta': 2 * np.arctan2(
                msg.pose.pose.orientation.z,
                msg.pose.pose.orientation.w),
            'linear_velocity': msg.twist.twist.linear.x,
            'angular_velocity': msg.twist.twist.angular.z,
            'timestamp': msg.header.stamp.sec
        }
        self.socketio.emit('robot_pose', self.robot_pose)

    def scan_callback(self, msg):
        """处理激光雷达数据"""
        self.laser_data = {
            'ranges': msg.ranges,
            'angle_min': msg.angle_min,
            'angle_max': msg.angle_max,
            'angle_increment': msg.angle_increment,
            'timestamp': msg.header.stamp.sec
        }
        self.socketio.emit('laser_scan', self.laser_data)

    def send_velocity_command(self, linear, angular):
        """发送速度命令"""
        twist = Twist()
        twist.linear.x = linear
        twist.angular.z = angular
        self.cmd_vel_pub.publish(twist)

    async def load_map(self, map_yaml_path):
        """加载新地图"""
        from nav2_msgs.srv import LoadMap
        
        req = LoadMap.Request()
        req.map_url = map_yaml_path
        
        await self.load_map_client.wait_for_service()
        future = self.load_map_client.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        
        return future.result()

    async def startup_nav2_stack(self):
        """启动Nav2堆栈"""
        from std_srvs.srv import Empty
        
        req = Empty.Request()
        await self.startup_nav2.wait_for_service()
        future = self.startup_nav2.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        
        return future.result()